import { Injectable, Logger, OnModuleInit, OnModuleDestroy } from '@nestjs/common';
import { ConfigService } from '@nestjs/config';
import * as WebSocket from 'ws';
import { EventEmitter } from 'events';
import * as net from 'net';
import fetch from 'node-fetch';
import { ComfyUIResponse, GenerationStatus } from '../types/comfyui.types';

interface TaskQueueItem {
  resolve: (value: any) => void;
  reject: (reason?: any) => void;
  timeout: NodeJS.Timeout;
}

interface TaskStatus {
  status: 'pending' | 'running' | 'completed' | 'failed';
  progress?: number;
  node?: string;
  output?: any;
  error?: string;
}

export class ComfyUIError extends Error {
  constructor(message: string, public code: string) {
    super(message);
    this.name = 'ComfyUIError';
  }
}

@Injectable()
export class ComfyUIService implements OnModuleInit, OnModuleDestroy {
  private readonly logger = new Logger(ComfyUIService.name);
  private readonly comfyuiUrl: string;
  private readonly wsUrl: string;
  private readonly requestTimeout = 300000;
  private readonly maxRetries = 3;
  private readonly retryDelay = 2000;
  private readonly maxReconnectAttempts = 5;
  private readonly reconnectInterval = 10000;
  private readonly pingTimeout = 10000;
  private connectionPromise: Promise<void> | null = null;

  private ws: WebSocket | null = null;
  private readonly events = new EventEmitter();
  private readonly taskQueue = new Map<string, TaskQueueItem>();
  private isConnected = false;
  private reconnectAttempts = 0;
  public client_id: string;

  private taskUpdateCallbacks = new Map<string, Function[]>();

  constructor(private configService: ConfigService) {
    const uri = this.configService.get<string>('COMFYUI_URI');
    if (!uri) {
      throw new Error('COMFYUI_URI environment variable is not set');
    }
    this.comfyuiUrl = `http://${uri}`;
    this.wsUrl = `ws://${uri}/ws`;

    // 生成唯一的 clientId
    this.client_id = `client_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`;

    // 打印配置信息
    this.logger.log('ComfyUI Configuration:', {
      httpUrl: this.comfyuiUrl,
      wsUrl: this.wsUrl
    });
  }

  async onModuleInit(): Promise<void> {
    // 初始化时开始后台重连，但不等待结果
    this.startBackgroundConnection();
  }

  onModuleDestroy(): void {
    this.disconnect();
  }

  private async startBackgroundConnection() {
    while (true) {
      try {
        if (!this.isConnected) {
          await this.connect();
        }
        // 连接成功，退出循环
        break;
      } catch (error) {
        this.logger.warn('Failed to connect to ComfyUI, will retry in background');
        await new Promise(resolve => setTimeout(resolve, this.reconnectInterval));
      }
    }
  }

  private async connect(): Promise<void> {
    if (this.isConnected) {
      this.logger.debug('Already connected to ComfyUI');
      return;
    }

    if (this.connectionPromise) {
      this.logger.debug('Connection attempt already in progress');
      return this.connectionPromise;
    }

    this.connectionPromise = (async () => {
      try {
        this.logger.debug('Starting connection process');

        // 检查服务可用性
        let isAvailable = false;
        for (let i = 0; i < 3; i++) {
          isAvailable = await this.pingComfyUI();
          if (isAvailable) break;
          if (i < 2) {
            await new Promise(resolve => setTimeout(resolve, 2000));
          }
        }

        if (!isAvailable) {
          throw new Error('ComfyUI service is not available');
        }


        this.logger.debug('Attempting WebSocket connection with clientId:', this.client_id);

        // 使用新的 WebSocket URL 格式
        this.ws = new WebSocket(`${this.wsUrl}?clientId=${this.client_id}`, {
          handshakeTimeout: this.pingTimeout,
          timeout: this.pingTimeout
        });

        await new Promise<void>((resolve, reject) => {
          const timeout = setTimeout(() => {
            this.logger.debug('WebSocket connection timeout');
            reject(new Error('WebSocket connection timeout'));
          }, this.pingTimeout);

          this.ws!.once('open', () => {
            this.logger.debug('WebSocket connection established successfully');
            clearTimeout(timeout);
            this.isConnected = true;
            this.reconnectAttempts = 0;
            this.setupWebSocketHandlers();
            resolve();
          });

          this.ws!.once('error', (error) => {
            clearTimeout(timeout);
            reject(error);
          });
        });

        this.logger.log('Successfully connected to ComfyUI WebSocket');

      } catch (error) {
        this.logger.error('Failed to connect:', error);
        throw error;
      } finally {
        this.connectionPromise = null;
      }
    })();

    return this.connectionPromise;
  }

  private disconnect(): void {
    if (this.ws) {
      this.ws.close();
      this.ws = null;
      this.isConnected = false;
    }
  }

  private async fetchWithRetry(url: string, options: RequestInit): Promise<Response> {
    let lastError: Error | null = null;

    for (let i = 0; i < this.maxRetries; i++) {
      try {
        const controller = new AbortController();
        const timeoutId = setTimeout(() => controller.abort(), this.requestTimeout);

        const response = await fetch(url, {
          ...options,
          signal: controller.signal,
        });

        clearTimeout(timeoutId);

        if (!response.ok) {
          throw new Error(`HTTP error! status: ${response.status}`);
        }

        return response;
      } catch (error) {
        lastError = error instanceof Error ? error : new Error('Unknown error');
        if (i < this.maxRetries - 1) {
          const delay = this.retryDelay * Math.pow(2, i);
          await new Promise(resolve => setTimeout(resolve, delay));
        }
      }
    }

    throw lastError || new Error('All retry attempts failed');
  }

  async executeWorkflow(data: any): Promise<ComfyUIResponse> {
    try {
      // 确保连接可用
      await this.ensureConnected();

      // ComfyUI API 格式
      const request = {
        client_id: this.client_id,
        prompt: data  // 工作流配置必须放在 prompt 字段下
      };

      // 提交任务到 ComfyUI
      const response = await fetch(`${this.comfyuiUrl}/prompt`, {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json'
        },
        body: JSON.stringify(request)
      });

      if (!response.ok) {
        const error = await response.json();
        throw new Error(`Failed to submit task: ${JSON.stringify(error)}`);
      }

      return await response.json();
    } catch (error) {
      this.logger.error('Workflow execution failed:', error);
      throw error;
    }
  }

  // 确保连接
  private async ensureConnected() {
    try {
      if (!this.isConnected) {
        this.logger.debug('Not connected, attempting to connect');
        await this.connect();
      }

      // 测试连接
      await this.testConnection();
      this.logger.debug('Connection test passed');
    } catch (error) {
      this.logger.error('Connection check failed:', error);
      throw new Error('无法连接到 ComfyUI 服务');
    }
  }

  // 测试连接
  private async testConnection(): Promise<void> {
    try {
      // 首先检查 TCP 连接
      const tcpConnected = await this.checkTcpConnection();
      if (!tcpConnected) {
        throw new Error('TCP connection failed');
      }
      this.logger.debug('TCP connection successful');

      // 然后检查 HTTP 连接
      this.logger.debug('Attempting basic HTTP connection');
      const response = await fetch(`${this.comfyuiUrl}/system_stats`);  // 使用 system_stats 端点进行测试
      if (!response.ok) {
        throw new Error(`HTTP check failed: ${response.status}`);
      }
      this.logger.debug('Basic HTTP connection successful');

    } catch (error) {
      this.logger.error('Connection test failed:', error);
      throw error;
    }
  }

  private async pingComfyUI(): Promise<boolean> {
    try {
      this.logger.debug('Starting ComfyUI availability check');

      const tcpCheck = await this.checkTcpConnection();
      if (!tcpCheck) {
        this.logger.warn('TCP connection failed - server might be down or unreachable');
        return false;
      }

      const controller = new AbortController();
      const timeoutId = setTimeout(() => {
        controller.abort();
        this.logger.debug('HTTP request aborted due to timeout');
      }, this.pingTimeout);

      try {
        this.logger.debug('Attempting basic HTTP connection');
        const response = await fetch(this.comfyuiUrl, {
          method: 'GET',
          signal: controller.signal,
        });

        if (response.ok) {
          this.logger.debug('Basic HTTP connection successful');
          clearTimeout(timeoutId);
          return true;
        }
      } catch (error) {
        this.logger.debug('Basic HTTP check failed:', error.message);
      }

      this.logger.debug('Attempting system_stats endpoint check');
      const response = await fetch(`${this.comfyuiUrl}/system_stats`, {
        method: 'GET',
        signal: controller.signal,
      });

      clearTimeout(timeoutId);
      const result = response.ok;
      this.logger.debug('System stats check result:', result);
      return result;
    } catch (error) {
      this.logger.error('Ping ComfyUI failed:', {
        error: error instanceof Error ? error.message : 'Unknown error',
        url: this.comfyuiUrl
      });
      return false;
    }
  }

  private async checkTcpConnection(): Promise<boolean> {
    try {
      const url = new URL(this.comfyuiUrl);
      this.logger.debug('Checking TCP connection:', {
        host: url.hostname,
        port: parseInt(url.port) || 80
      });

      const net = require('net');

      return new Promise((resolve) => {
        const socket = new net.Socket();

        socket.setTimeout(5000);

        socket.on('connect', () => {
          this.logger.debug('TCP connection successful');
          socket.destroy();
          resolve(true);
        });

        socket.on('error', (error) => {
          this.logger.debug('TCP connection error:', error.message);
          socket.destroy();
          resolve(false);
        });

        socket.on('timeout', () => {
          this.logger.debug('TCP connection timeout');
          socket.destroy();
          resolve(false);
        });

        socket.connect(parseInt(url.port) || 80, url.hostname);
      });
    } catch (error) {
      this.logger.error('TCP connection check failed:', error);
      return false;
    }
  }

  private setupWebSocketHandlers(): void {
    if (!this.ws) return;

    this.ws.on('message', (data: WebSocket.Data) => {
      try {
        const message = JSON.parse(data.toString());
        this.handleComfyUIWebSocketMessage(message);
      } catch (error) {
        this.logger.error('Failed to parse WebSocket message:', error);
      }
    });

    this.ws.on('close', () => {
      this.logger.warn('ComfyUI WebSocket connection closed');
      this.handleDisconnect();
    });

    this.ws.on('error', (error) => {
      this.logger.error('ComfyUI WebSocket error:', error);
      this.handleDisconnect();
    });
  }

  private handleComfyUIWebSocketMessage(message: any) {
    // console.log("Backend ComfyUI Message:", message);

    if (message.type === 'execution_start') {
      this.notifyTaskUpdate(message.data.prompt_id, {
        type: 'execution_start',
        progress: 0,
        comfyui_data: message.data  // 传递原始数据
      });
    }
    else if (message.type === 'progress') {
      this.notifyTaskUpdate(message.data.prompt_id, {
        type: 'progress',
        progress: message.data.value || 0,
        value: message.data.value || 0,
        max: message.data.max || 0,
        comfyui_data: message.data  // 包含 value, max, prompt_id, node 等信息
      });
    }
    else if (message.type === 'executing') {
      this.notifyTaskUpdate(message.data.prompt_id, {
        type: 'executing',
        progress: message.data.value || 0,
        comfyui_data: message.data  // 包含 value, max, prompt_id, node 等信息
      });
    }
    else if (message.type === 'executed') {
      this.notifyTaskUpdate(message.data.prompt_id, {
        type: 'executed',
        progress: message.data.value || 0,
        comfyui_data: message.data  // 包含 value, max, prompt_id, node 等信息
      });
    }
    else if (message.type === 'execution_cached') {
      this.notifyTaskUpdate(message.data.prompt_id, {
        type: 'execution_cached',
        progress: 100,
        comfyui_data: message.data
      });
    }
    else if (message.type === 'execution_success') {
      this.notifyTaskUpdate(message.data.prompt_id, {
        type: 'execution_success',
        outputs: message.data.output,
        comfyui_data: message.data
      });
    }
    else if (message.type === 'execution_error') {
      this.notifyTaskUpdate(message.data.prompt_id, {
        type: 'error',
        error: message.data.error || '生成失败',
        comfyui_data: message.data
      });
    }
  }

  private handleDisconnect(): void {
    this.isConnected = false;
    this.ws = null;

    if (this.reconnectAttempts < this.maxReconnectAttempts) {
      this.reconnectAttempts++;
      this.logger.log(`Will attempt to reconnect (${this.reconnectAttempts}/${this.maxReconnectAttempts}) in ${this.reconnectInterval}ms`);
      // 启动后台重连
      setTimeout(() => this.startBackgroundConnection(), this.reconnectInterval);
    } else {
      this.logger.warn('Max reconnection attempts reached');
      this.taskQueue.forEach(({ reject, timeout }) => {
        clearTimeout(timeout);
        reject(new Error('Connection lost'));
      });
      this.taskQueue.clear();
      // 重置重连次数，允许之后再次尝试
      this.reconnectAttempts = 0;
    }
  }

  // 添加状态监听方法
  onStatus(callback: (status: any) => void): () => void {
    this.events.on('status', callback);
    return () => this.events.off('status', callback);
  }

  // 添加进度监听方法
  onProgress(callback: (progress: any) => void): () => void {
    this.events.on('progress', callback);
    return () => this.events.off('progress', callback);
  }

  // 修改状态检查方法
  public async checkTaskStatus(promptId: string): Promise<any> {
    try {
      this.logger.debug('Checking task status:', promptId);

      // 1. 检查队列状态
      const queueResponse = await this.fetchWithRetry(`${this.comfyuiUrl}/queue`, {
        method: 'GET',
        headers: {
          'Content-Type': 'application/json'
        }
      });

      const queueData = await queueResponse.json();

      // 检查运行中的任务
      const runningItem = queueData.queue_running?.find((item: any) =>
        item.prompt_id === promptId ||
        item.prompt?.prompt_id === promptId
      );

      if (runningItem) {
        return {
          status: 'running',
          progress: runningItem.progress || 0,
          node: runningItem.node_id
        };
      }

      // 检查等待中的任务
      const pendingItem = queueData.queue_pending?.find((item: any) =>
        item.prompt_id === promptId ||
        item.prompt?.prompt_id === promptId
      );

      if (pendingItem) {
        return {
          status: 'pending',
          progress: 0,
          node: null
        };
      }

      // 2. 检查历史记录
      try {
        const historyResponse = await this.fetchWithRetry(`${this.comfyuiUrl}/history/${promptId}`, {
          method: 'GET',
          headers: {
            'Content-Type': 'application/json'
          }
        });

        const history = await historyResponse.json();

        if (history && Object.keys(history).length > 0) {
          // 检查是否有输出
          const outputs = history[promptId]?.outputs;
          if (outputs && Object.keys(outputs).length > 0) {
            return {
              status: 'completed',
              outputs: outputs,
              progress: 100
            };
          }
        }
      } catch (historyError) {
        this.logger.warn('History check failed:', historyError);
        // 历史记录检查失败不应该立即返回失败状态
      }

      // 3. 如果任务刚刚提交，给它一些时间进入队列
      return {
        status: 'pending',
        progress: 0,
        message: '等待任务开始'
      };

    } catch (error) {
      this.logger.error('Check task status failed:', error);

      // 如果是网络错误，返回pending状态而不是失败
      if (error instanceof Error && error.message.includes('fetch')) {
        return {
          status: 'pending',
          progress: 0,
          message: '正在连接服务器'
        };
      }

      throw new Error(`检查任务状态失败: ${error instanceof Error ? error.message : '未知错误'}`);
    }
  }

  // 修改历史记录获取方法
  async getHistory(promptId: string): Promise<any> {
    try {
      this.logger.debug('Getting history for prompt:', promptId);

      const response = await this.fetchWithRetry(`${this.comfyuiUrl}/history/${promptId}`, {
        method: 'GET',
        headers: {
          'Content-Type': 'application/json'
        }
      });

      const history = await response.json();

      // 直接返回特定promptId的历史记录
      if (history && history[promptId]) {
        return history[promptId];
      }

      return null;

    } catch (error) {
      this.logger.warn('Get history failed:', error);
      // 获取历史记录失败不应该抛出错误，而是返回null
      return null;
    }
  }

  // 添加一个辅助方法用于检查任务是否完成
  private isTaskCompleted(outputs: any): boolean {
    if (!outputs) return false;

    // 检查是否所有输出节点都有图片
    return Object.values(outputs).some((output: any) =>
      output.images && output.images.length > 0
    );
  }

  onTaskUpdate(promptId: string, callback: (data: any) => void) {
    if (!this.taskUpdateCallbacks.has(promptId)) {
      this.taskUpdateCallbacks.set(promptId, []);
    }
    this.taskUpdateCallbacks.get(promptId)?.push(callback);
  }

  private notifyTaskUpdate(promptId: string, data: any) {
    const callbacks = this.taskUpdateCallbacks.get(promptId);
    if (callbacks) {
      callbacks.forEach(callback => callback(data));
    }
  }

  private async handleExecutionComplete(promptId: string, output: any): Promise<void> {
    try {
      this.logger.debug(`Execution complete for prompt: ${promptId}`);

      // 记录完整的输出以便调试
      this.logger.debug(`Raw output: ${JSON.stringify(output, null, 2)}`);

      // 添加对图片路径的处理逻辑
      const processedOutput = this.processOutputImages(output);

      // 通知任务更新
      this.notifyTaskUpdate(promptId, {
        type: 'completed',
        outputs: processedOutput,
        comfyui_data: output  // 也传递原始数据供前端参考
      });
    } catch (error) {
      this.logger.error(`Error handling execution complete: ${error.message}`);
    }
  }

  // 添加处理图片路径的辅助方法
  private processOutputImages(output: any): any {
    if (!output) return output;

    // 深拷贝以避免修改原始数据
    const processed = JSON.parse(JSON.stringify(output));

    // 遍历所有节点输出
    for (const nodeId in processed) {
      const nodeOutput = processed[nodeId];
      if (nodeOutput && nodeOutput.images && Array.isArray(nodeOutput.images)) {
        // 确保图片的filename是完整路径
        nodeOutput.images = nodeOutput.images.map(img => {
          if (img.filename && !img.filename.startsWith('/')) {
            // 确保路径格式与前端期望一致
            this.logger.debug(`Processing image path: ${img.filename}`);
          }
          return img;
        });
      }
    }

    return processed;
  }
}